from matplotlib import pyplot as plt 
import torch 
import torch.nn as nn  # 添加此行导入 torch.nn
import pandas as pd

# 获取数据，使用pandas库
data = pd.read_csv('./dateset/Income1.csv')

# 绘制散点图
plt.scatter(data.Education, data.Income)
# 设置x轴,y轴坐标
plt.xlabel('Education')
plt.ylabel('Income')

# pandas数据转为pytorch数据（张量）
X = torch.from_numpy(data.Education.to_numpy().reshape(-1,1)).float()  # 转换为浮点类型
Y = torch.from_numpy(data.Income.to_numpy().reshape(-1,1)).float()  # 转换为浮点类型

# 导入用于搭建神经网络模型
# class EasyModel(nn.Module): 
#     # 实现一层线性网络
#     def __init__(self): 
#         super(EasyModel, self).__init__()
#         # 定义中间线性层，输入维度为1，输出维度为1
#         self.linear = nn.Linear(in_features=1, out_features=1)
#     def forward(self, x): 
#         # 定义前向传播
#         out = self.linear(x)
#         return out

class EasyModel(nn.Module):
    #初始化模型，实现1层线性网络
    def __init__(self):
        #调用父类进行初始化
        super(EasyModel,self).__init__()
        #定义中间的一层线性层，输入的维度是1，输出的维度也是1
        self.linear = nn.Linear(in_features=1,out_features=1)
        
    #定义1个前向传播函数，告诉模型如何一层一层往下输出的逻辑
    def forward(self,inputs):
        #设置将输入的教育支出，输入到线形层
        logits = self.linear(inputs)
        return logits


model = EasyModel()

#定义均方误差为损失函数
loss_fn = nn.MSELoss()

# 定义优化器（怎样让损失函数最小）
# todo 学习率0.01时报错？
opt = torch.optim.SGD(model.parameters(),lr=0.0001) # 优化器为随机梯度下降，学习率0.0001,调节学习速度快慢

# 开始训练
for epoch in range(5000):
    for x,y in zip(X,Y):
        #先让模型做预测
        y_pred = model(x)
        # 预测了，看看误差咋样，计算损失函数
        loss = loss_fn(y,y_pred)
        # 复原，梯度归零
        opt.zero_grad()
        # 根据误差，倒推应该怎么修改参数，反向传播
        loss.backward()
        # 使用优化器，更新参数
        opt.step()

# 预测结果 detach:用于前项传播计算
Y_pred = model(X).detach()
Y_pred = Y_pred.numpy()

print(Y_pred)

# 再绘制散点图
plt.scatter(data.Education, data.Income, color ='red', label = 'real')
plt.scatter(data.Education, Y_pred, c='g', label = 'pred')
# 设置x轴,y轴坐标
plt.xlabel('Education')
plt.ylabel('Income')

plt.savefig('./dateset/income.png')